# -*- coding: utf-8 -*-
from echo.reservoir import Reservoir
from mlp.MLPModel import MLPModel
from data_parse.ParsePrepareDataEchoNet import DataParse
from utils import utils
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import const
import time
import sys

def transfromModelToLite():
	try:
		modelDir = str(sys.argv[2])
	except:
		modelDir = const.PATH_TO_DIR

	mlpModel = MLPModel(loadFromConst=True, modelDir=modelDir)
	mlpModel.convertToTfLite(modelDir)

def manualTraining():
	maxModelValue = 0
	largeEpoch = 10000
	while largeEpoch:
		isEchoLoad 	 = bool(int(sys.argv[2]))
		isMLPLoad 	 = bool(int(sys.argv[3]))
		echoDir 	 = str(sys.argv[4])
		mlpDir 		 = str(sys.argv[5])
		useMlpDiydiy = str(sys.argv[6])
		# 创建网络
		echo = Reservoir(loadFromConst = isEchoLoad, modelDir = echoDir)
		mlpModel = MLPModel(loadFromConst = isMLPLoad, modelDir = mlpDir, diy=useMlpDiydiy)
		# 解析数据
		dataParse = DataParse()
		_, x, y, _, _, _ = dataParse.generateTrainAndTestData(const.READ_DATA_ARRAY_TYPE, 1)
		# 数据增强（平均采样）
		xAfterAug, yAfterAug = dataParse.dataAugmentation(x, y)
		# ground true 预处理
		trainYClassic = utils.oneHotEncodingToClassicEncoding(yAfterAug)
		# 模型训练
		xAfterEcho   = echo.getStates(xAfterAug, 0, False)
		mlpModel.trainModel(xAfterEcho, trainYClassic, 0.1)
		# 分析保存结果
		a = np.array(mlpModel.getLastOneHistoryOnFloat())
		curModelVale = (a * np.array([2, 10, 1, 10])).sum()
		if maxModelValue < curModelVale:
			saveAll(echo, mlpModel) 
			maxModelValue = curModelVale

		largeEpoch -= 1

def saveAll(echo, mlpModel):
	print ('save the model ...')
	savePath = './model-%s-%s/'%(time.strftime("%Y-%m-%d", time.localtime()), mlpModel.getLastOneHistoryOnStr())
	mlpModel.saveModel(savePath)
	mlpModel.saveHistory(savePath)
	echo.saveModel(savePath)

def testOffLine():
	# 加载网络
	echo = Reservoir(loadFromConst = True)
	mlpModel = MLPModel(loadFromConst = True)
	df = pd.read_csv('./off_line_data.csv').values
	length, _ = df.shape
	df = df.tolist()
	# 加载数据
	data = []
	for i in range(0, length):
		sample = [df[j] for j in range(i, i+const.T*const.AUG_STEPS, const.AUG_STEPS) if j < length]
		if len(sample) == const.T:
			data.append(sample)
	data = np.array(data)
	# 计算
	xAfterEcho = echo.getStates(data, 0, False)
	yPredictOneHot = mlpModel.calcuModel(xAfterEcho)
	yPredictClaissic = utils.oneHotEncodingToClassicEncoding(yPredictOneHot)
	# 输出结果
	plotForTwoDimArray(yPredictOneHot)
	resStr = ''
	lastChar = ''
	num = 0
	for i, classic in enumerate(yPredictClaissic):
		c = const.GESTURES_ICON[classic]
		if c == '_' and lastChar == '_' or c == lastChar:
			if lastChar == c:
				num += 1
			continue
		if c != lastChar:
			resStr += '%s*%s\n'%(c, num)
			lastChar = c
			num = 0
	print (resStr)

def plotForTwoDimArray(d):
	index = range(1, len(d) + 1)
	color = ['r', 'g', 'b', 'y', 'c', 'm', 'k', 'aquamarine', 'dodgerblue']
	for i in range(0, len(d[0])):
		plt.scatter(index, d[:,i], color=color[i], label=('%s'%const.GESTURES_ICON[i]), marker='.', s=16, alpha=0.5)
	plt.legend()
	plt.show()

def main():
	arg1 = int(sys.argv[1])
	if arg1 == 1:
		manualTraining()
	if arg1 == 2:
		transfromModelToLite()
	if arg1 == 3:
		testOffLine()

if __name__ == '__main__':
	print ("1 : 手动训练神经网络，参数[是否加载echo网络 , 是否加载mlp模型 , echo网络目录地址 , mlp模型目录地址 , mlp是否diy]")
	print ("2 : 将mlp转换为轻量级模型 , 参数[mlp模型目录地址]")
	print ("3 : 测试离线数据")
	if len(sys.argv) > 1:
		main()